package net.cyue.ort.llm.data;

import java.util.*;
import java.util.stream.Collectors;

/**
 * Token集合类
 * 提供token的唯一集合管理
 */
public class TokenSet implements Collection<Token>, AutoCloseable {
    private final Set<Token> tokenSet;
    
    public TokenSet() {
        this.tokenSet = new HashSet<>();
    }
    
    public TokenSet(Set<Token> tokenSet) {
        this.tokenSet = new HashSet<>(tokenSet);
    }
    
    @Override
    public void close() {
        tokenSet.clear();
    }
    
    /**
     * 获取所有token ID数组
     */
    public long[] getTokenIDs() {
        return tokenSet.stream().mapToLong(Token::getId).toArray();
    }
    
    /**
     * 获取所有token ID列表
     */
    public List<Long> getTokenIDList() {
        return tokenSet.stream().map(Token::getId).collect(Collectors.toList());
    }
    
    /**
     * 检查是否包含指定token
     */
    public boolean contains(Token token) {
        return tokenSet.contains(token);
    }
    
    /**
     * 检查是否包含指定token ID
     */
    public boolean contains(long tokenID) {
        return tokenSet.stream().anyMatch(token -> token.getId() == tokenID);
    }
    
    /**
     * 添加token
     */
    @Override
    public boolean add(Token token) {
        return tokenSet.add(token);
    }
    
    /**
     * 添加token（别名方法）
     */
    public void addToken(Token token) {
        tokenSet.add(token);
    }
    
    /**
     * 添加TokenSet中的所有token
     */
    public void addAll(TokenSet tokenSet) {
        this.tokenSet.addAll(tokenSet.tokenSet);
    }
    
    /**
     * 添加TokenList中的所有token
     */
    public void addAll(TokenList tokenList) {
        this.tokenSet.addAll(tokenList);
    }
    
    /**
     * 添加Set中的所有token
     */
    public void addAll(Set<Token> tokenSet) {
        this.tokenSet.addAll(tokenSet);
    }
    
    /**
     * 添加List中的所有token
     */
    public void addAll(List<Token> tokenList) {
        this.tokenSet.addAll(tokenList);
    }
    
    /**
     * 添加TokenSet（别名方法）
     */
    public void addTokenSet(TokenSet tokenSet) {
        addAll(tokenSet);
    }
    
    /**
     * 添加TokenList（别名方法）
     */
    public void addTokenList(TokenList tokenList) {
        addAll(tokenList);
    }
    
    /**
     * 移除指定token
     */
    public void remove(Token token) {
        tokenSet.remove(token);
    }
    
    /**
     * 移除指定token ID的token
     */
    public void remove(long tokenID) {
        tokenSet.removeIf(token -> token.getId() == tokenID);
    }
    
    /**
     * 清空集合
     */
    public void clear() {
        tokenSet.clear();
    }
    
    // Collection接口实现
    
    @Override
    public int size() {
        return tokenSet.size();
    }
    
    @Override
    public boolean isEmpty() {
        return tokenSet.isEmpty();
    }
    
    @Override
    public boolean contains(Object o) {
        return tokenSet.contains(o);
    }
    
    @Override
    public Iterator<Token> iterator() {
        return tokenSet.iterator();
    }
    
    @Override
    public Object[] toArray() {
        return tokenSet.toArray();
    }
    
    @Override
    public <T> T[] toArray(T[] a) {
        return tokenSet.toArray(a);
    }
    
    @Override
    public boolean remove(Object o) {
        return tokenSet.remove(o);
    }
    
    @Override
    public boolean containsAll(Collection<?> c) {
        return tokenSet.containsAll(c);
    }
    
    @Override
    public boolean addAll(Collection<? extends Token> c) {
        return tokenSet.addAll(c);
    }
    
    @Override
    public boolean removeAll(Collection<?> c) {
        return tokenSet.removeAll(c);
    }
    
    @Override
    public boolean retainAll(Collection<?> c) {
        return tokenSet.retainAll(c);
    }
}

